package generatekey

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"io/ioutil"
	"os"
)

// 生成公钥和私钥,bitSize为密钥的长度,建议使用2048
func MakeKey(bitSize int, privateKeyFilePath string, publicKeyFilePath string) {
	privateKey, err := rsa.GenerateKey(rand.Reader, bitSize)
	if err != nil {
		fmt.Println(err)
	}
	privateKeyBlock := &pem.Block{
		Type:    "RSA PRIVATE KEY",
		Headers: nil,
		Bytes:   x509.MarshalPKCS1PrivateKey(privateKey),
	}

	privateKeyfile, _ := os.Create(privateKeyFilePath)
	defer privateKeyfile.Close()
	pem.Encode(privateKeyfile, privateKeyBlock)

	publicKeyBlock := &pem.Block{
		Type:    "RSA PUBLIC KEY",
		Headers: nil,
		Bytes:   x509.MarshalPKCS1PublicKey(&privateKey.PublicKey),
	}
	publicKeyfile, _ := os.Create(publicKeyFilePath)
	defer publicKeyfile.Close()
	pem.Encode(publicKeyfile, publicKeyBlock)
}

// 根据私钥生成对应的公钥
func GeneratePublibKey(privateKeyFilePath string, publicKeyFilePath string) {
	prikeyfile, err := os.Open(privateKeyFilePath)
	if err != nil {
		fmt.Println(err.Error())
		return
	}
	// 读取私钥文件内容
	prikeyfilebytes, err := ioutil.ReadAll(prikeyfile)
	if len(prikeyfilebytes) >= 0 {
		if err != nil {
			fmt.Println(err.Error())
			return
		}
	} else {
		fmt.Println("This private key file is NULL!")
		return
	}
	// 将私钥文件转换成pem格式的对象
	prikeyblock, _ := pem.Decode(prikeyfilebytes)
	if prikeyblock == nil || prikeyblock.Type != "RSA PRIVATE KEY" {
		fmt.Println("failed to decode PEM block containing private key")
		return
	}
	// 提取私钥值
	privateKey, err := x509.ParsePKCS1PrivateKey(prikeyblock.Bytes)
	if err != nil {
		fmt.Println(err.Error())
		return
	}

	publicKeyBlock := &pem.Block{
		Type:    "RSA PUBLIC KEY",
		Headers: nil,
		Bytes:   x509.MarshalPKCS1PublicKey(&privateKey.PublicKey),
	}
	publicKeyfile, _ := os.Create(publicKeyFilePath)
	defer publicKeyfile.Close()
	pem.Encode(publicKeyfile, publicKeyBlock)

}
